#!/bin/bash

# 设置颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 默认参数
INTERVAL=5      # 采样间隔（秒）
DURATION=60     # 监控时长（秒）
OUTPUT_FILE="storage_performance_report_$(date +%Y%m%d_%H%M%S).txt"

# 帮助信息
show_help() {
    echo "用法: $0 [选项]"
    echo "选项:"
    echo "  -i, --interval <秒>   采样间隔（默认: 5秒）"
    echo "  -d, --duration <秒>   监控时长（默认: 60秒）"
    echo "  -o, --output <文件>   输出文件（默认: storage_performance_report_时间戳.txt）"
    echo "  -h, --help           显示此帮助信息"
    exit 0
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -i|--interval)
            INTERVAL="$2"
            shift 2
            ;;
        -d|--duration)
            DURATION="$2"
            shift 2
            ;;
        -o|--output)
            OUTPUT_FILE="$2"
            shift 2
            ;;
        -h|--help)
            show_help
            ;;
        *)
            echo "未知选项: $1"
            show_help
            ;;
    esac
done

# 检查必要的命令
check_commands() {
    local missing_commands=()
    
    for cmd in iostat lsblk df; do
        if ! command -v $cmd &> /dev/null; then
            missing_commands+=($cmd)
        fi
    done
    
    if [ ${#missing_commands[@]} -ne 0 ]; then
        echo -e "${RED}错误: 以下命令不可用: ${missing_commands[*]}${NC}"
        echo "请安装必要的工具包:"
        echo "  - iostat: sysstat 包"
        echo "  - lsblk: util-linux 包"
        echo "  - df: coreutils 包"
        exit 1
    fi
}

# 创建输出文件
echo "=== Linux 存储性能监控报告 ===" > "$OUTPUT_FILE"
echo "开始时间: $(date)" >> "$OUTPUT_FILE"
echo "采样间隔: ${INTERVAL}秒" >> "$OUTPUT_FILE"
echo "监控时长: ${DURATION}秒" >> "$OUTPUT_FILE"
echo "----------------------------------------" >> "$OUTPUT_FILE"

# 检查命令
check_commands

# 初始化计数器
count=0
total_samples=$((DURATION / INTERVAL))

# 创建临时文件存储数据
temp_disk_io="/tmp/disk_io_$$.tmp"
temp_disk_util="/tmp/disk_util_$$.tmp"
temp_disk_latency="/tmp/disk_latency_$$.tmp"
temp_disk_throughput="/tmp/disk_throughput_$$.tmp"
temp_disk_errors="/tmp/disk_errors_$$.tmp"
temp_disk_queue="/tmp/disk_queue_$$.tmp"
temp_filesystem_usage="/tmp/filesystem_usage_$$.tmp"
temp_io_wait="/tmp/io_wait_$$.tmp"

# 创建临时文件
touch "$temp_disk_io" "$temp_disk_util" "$temp_disk_latency" "$temp_disk_throughput" \
      "$temp_disk_errors" "$temp_disk_queue" "$temp_filesystem_usage" "$temp_io_wait"

# 清理函数
cleanup() {
    # 检查是否已经清理过
    if [ -n "$CLEANUP_DONE" ]; then
        return
    fi
    
    rm -f "$temp_disk_io" "$temp_disk_util" "$temp_disk_latency" "$temp_disk_throughput" \
          "$temp_disk_errors" "$temp_disk_queue" "$temp_filesystem_usage" "$temp_io_wait"
    echo -e "\n${GREEN}监控已完成，报告已保存到: $OUTPUT_FILE${NC}"
    
    # 标记清理已完成
    CLEANUP_DONE=1
    
    # 只有在非正常退出时才退出
    if [ -n "$1" ]; then
        exit 1
    fi
}

# 设置清理陷阱
trap 'cleanup 1' EXIT INT TERM

echo -e "${BLUE}开始监控存储性能...${NC}"
echo -e "采样间隔: ${INTERVAL}秒"
echo -e "监控时长: ${DURATION}秒"
echo -e "输出文件: $OUTPUT_FILE"
echo -e "按 Ctrl+C 可以提前结束监控\n"

# 获取所有块设备
get_block_devices() {
    lsblk -d -n -o NAME | grep -v "loop"
}

# 主循环
while [ $count -lt $total_samples ]; do
    timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    # 收集磁盘I/O统计
    iostat -x 1 1 | while read -r line; do
        if [[ $line =~ ^[a-zA-Z0-9/]+[[:space:]]+[0-9.]+ ]]; then
            device=$(echo "$line" | awk '{print $1}')
            if [[ $device != "Device" && $device != "loop"* ]]; then
                rrqm=$(echo "$line" | awk '{print $2}')
                wrqm=$(echo "$line" | awk '{print $3}')
                r=$(echo "$line" | awk '{print $4}')
                w=$(echo "$line" | awk '{print $5}')
                rsec=$(echo "$line" | awk '{print $6}')
                wsec=$(echo "$line" | awk '{print $7}')
                echo "$timestamp DISK_IO $device $rrqm $wrqm $r $w $rsec $wsec" >> "$temp_disk_io"
            fi
        fi
    done
    
    # 收集磁盘利用率
    iostat -x 1 1 | while read -r line; do
        if [[ $line =~ ^[a-zA-Z0-9/]+[[:space:]]+[0-9.]+ ]]; then
            device=$(echo "$line" | awk '{print $1}')
            if [[ $device != "Device" && $device != "loop"* ]]; then
                util=$(echo "$line" | awk '{print $NF}')
                echo "$timestamp DISK_UTIL $device $util" >> "$temp_disk_util"
            fi
        fi
    done
    
    # 收集磁盘延迟
    iostat -x 1 1 | while read -r line; do
        if [[ $line =~ ^[a-zA-Z0-9/]+[[:space:]]+[0-9.]+ ]]; then
            device=$(echo "$line" | awk '{print $1}')
            if [[ $device != "Device" && $device != "loop"* ]]; then
                await=$(echo "$line" | awk '{print $10}')
                svctm=$(echo "$line" | awk '{print $11}')
                echo "$timestamp DISK_LATENCY $device $await $svctm" >> "$temp_disk_latency"
            fi
        fi
    done
    
    # 收集磁盘吞吐量
    iostat -d 1 1 | while read -r line; do
        if [[ $line =~ ^[a-zA-Z0-9/]+[[:space:]]+[0-9.]+ ]]; then
            device=$(echo "$line" | awk '{print $1}')
            if [[ $device != "Device" && $device != "loop"* ]]; then
                tps=$(echo "$line" | awk '{print $2}')
                kb_read=$(echo "$line" | awk '{print $3}')
                kb_wrtn=$(echo "$line" | awk '{print $4}')
                echo "$timestamp DISK_THROUGHPUT $device $tps $kb_read $kb_wrtn" >> "$temp_disk_throughput"
            fi
        fi
    done
    
    # 收集磁盘错误
    for device in $(get_block_devices); do
        if [ -f "/sys/block/$device/stat" ]; then
            read -r io sectors ticks ios sectors_discard ios_discard < "/sys/block/$device/stat"
            if [ -f "/sys/block/$device/queue/hw_sector_size" ]; then
                sector_size=$(cat "/sys/block/$device/queue/hw_sector_size")
            else
                sector_size=512
            fi
            echo "$timestamp DISK_ERRORS $device $io $sectors $ticks $ios $sectors_discard $ios_discard $sector_size" >> "$temp_disk_errors"
        fi
    done
    
    # 收集磁盘队列信息
    for device in $(get_block_devices); do
        if [ -f "/sys/block/$device/queue/nr_requests" ]; then
            nr_requests=$(cat "/sys/block/$device/queue/nr_requests")
            if [ -f "/sys/block/$device/queue/nr_in_flight" ]; then
                nr_in_flight=$(cat "/sys/block/$device/queue/nr_in_flight")
            else
                nr_in_flight=0
            fi
            echo "$timestamp DISK_QUEUE $device $nr_requests $nr_in_flight" >> "$temp_disk_queue"
        fi
    done
    
    # 收集文件系统使用情况
    df -h | while read -r line; do
        if [[ $line =~ ^/dev/ ]]; then
            device=$(echo "$line" | awk '{print $1}')
            size=$(echo "$line" | awk '{print $2}')
            used=$(echo "$line" | awk '{print $3}')
            avail=$(echo "$line" | awk '{print $4}')
            usage=$(echo "$line" | awk '{print $5}')
            mount=$(echo "$line" | awk '{print $6}')
            echo "$timestamp FILESYSTEM_USAGE $device $size $used $avail $usage $mount" >> "$temp_filesystem_usage"
        fi
    done
    
    # 收集I/O等待时间
    vmstat 1 1 | tail -n 1 | while read -r line; do
        io_wait=$(echo "$line" | awk '{print $16}')
        echo "$timestamp IO_WAIT $io_wait" >> "$temp_io_wait"
    done
    
    # 显示倒计时（每秒更新一次）
    remaining=$((DURATION - (count * INTERVAL)))
    for ((i=0; i<INTERVAL; i++)); do
        printf "\r${YELLOW}剩余时间: %d 秒${NC}" $((remaining - i))
        sleep 1
    done
    
    count=$((count + 1))
done

echo -e "\n\n${GREEN}正在生成报告...${NC}"

# 生成统计报告
echo -e "\n=== 磁盘I/O统计 ===" >> "$OUTPUT_FILE"
echo "时间戳                设备    合并读请求  合并写请求  读请求/秒  写请求/秒  读扇区/秒  写扇区/秒" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"
current_timestamp=""
while read -r line; do
    if [[ "$line" == *"DISK_IO"* ]]; then
        timestamp=$(echo "$line" | cut -d' ' -f1,2)
        device=$(echo "$line" | cut -d' ' -f4)
        rrqm=$(echo "$line" | cut -d' ' -f5)
        wrqm=$(echo "$line" | cut -d' ' -f6)
        r=$(echo "$line" | cut -d' ' -f7)
        w=$(echo "$line" | cut -d' ' -f8)
        rsec=$(echo "$line" | cut -d' ' -f9)
        wsec=$(echo "$line" | cut -d' ' -f10)
        
        # 如果时间戳变化，添加一个空行
        if [[ -n "$current_timestamp" && "$current_timestamp" != "$timestamp" ]]; then
            echo "" >> "$OUTPUT_FILE"
        fi
        
        awk -v ts="$timestamp" -v d="$device" -v rr="$rrqm" -v wr="$wrqm" -v rd="$r" -v wd="$w" -v rs="$rsec" -v ws="$wsec" \
            'BEGIN {printf "%-20s %-8s %10.2f %10.2f %10.2f %10.2f %10.2f %10.2f\n", ts, d, rr, wr, rd, wd, rs, ws}' >> "$OUTPUT_FILE"
        
        current_timestamp="$timestamp"
    fi
done < "$temp_disk_io"

echo -e "\n=== 磁盘利用率统计 ===" >> "$OUTPUT_FILE"
echo "时间戳                设备    利用率(%)" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"
current_timestamp=""
while read -r line; do
    if [[ "$line" == *"DISK_UTIL"* ]]; then
        timestamp=$(echo "$line" | cut -d' ' -f1,2)
        device=$(echo "$line" | cut -d' ' -f4)
        util=$(echo "$line" | cut -d' ' -f5)
        
        # 如果时间戳变化，添加一个空行
        if [[ -n "$current_timestamp" && "$current_timestamp" != "$timestamp" ]]; then
            echo "" >> "$OUTPUT_FILE"
        fi
        
        awk -v ts="$timestamp" -v d="$device" -v u="$util" \
            'BEGIN {printf "%-20s %-8s %10.2f%%\n", ts, d, u}' >> "$OUTPUT_FILE"
        
        current_timestamp="$timestamp"
    fi
done < "$temp_disk_util"

echo -e "\n=== 磁盘延迟统计 ===" >> "$OUTPUT_FILE"
echo "时间戳                设备    平均等待时间(ms)  服务时间(ms)" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"
current_timestamp=""
while read -r line; do
    if [[ "$line" == *"DISK_LATENCY"* ]]; then
        timestamp=$(echo "$line" | cut -d' ' -f1,2)
        device=$(echo "$line" | cut -d' ' -f4)
        await=$(echo "$line" | cut -d' ' -f5)
        svctm=$(echo "$line" | cut -d' ' -f6)
        
        # 如果时间戳变化，添加一个空行
        if [[ -n "$current_timestamp" && "$current_timestamp" != "$timestamp" ]]; then
            echo "" >> "$OUTPUT_FILE"
        fi
        
        awk -v ts="$timestamp" -v d="$device" -v a="$await" -v s="$svctm" \
            'BEGIN {printf "%-20s %-8s %15.2f %15.2f\n", ts, d, a, s}' >> "$OUTPUT_FILE"
        
        current_timestamp="$timestamp"
    fi
done < "$temp_disk_latency"

echo -e "\n=== 磁盘吞吐量统计 ===" >> "$OUTPUT_FILE"
echo "时间戳                设备    每秒传输量  读取(KB/s)  写入(KB/s)" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"
current_timestamp=""
while read -r line; do
    if [[ "$line" == *"DISK_THROUGHPUT"* ]]; then
        timestamp=$(echo "$line" | cut -d' ' -f1,2)
        device=$(echo "$line" | cut -d' ' -f4)
        tps=$(echo "$line" | cut -d' ' -f5)
        kb_read=$(echo "$line" | cut -d' ' -f6)
        kb_wrtn=$(echo "$line" | cut -d' ' -f7)
        
        # 如果时间戳变化，添加一个空行
        if [[ -n "$current_timestamp" && "$current_timestamp" != "$timestamp" ]]; then
            echo "" >> "$OUTPUT_FILE"
        fi
        
        awk -v ts="$timestamp" -v d="$device" -v t="$tps" -v kr="$kb_read" -v kw="$kb_wrtn" \
            'BEGIN {printf "%-20s %-8s %10.2f %10.2f %10.2f\n", ts, d, t, kr, kw}' >> "$OUTPUT_FILE"
        
        current_timestamp="$timestamp"
    fi
done < "$temp_disk_throughput"

echo -e "\n=== 磁盘错误统计 ===" >> "$OUTPUT_FILE"
echo "时间戳                设备    读操作数  读扇区数  读时间(ms)  写操作数  写扇区数  写时间(ms)  扇区大小" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"
current_timestamp=""
while read -r line; do
    if [[ "$line" == *"DISK_ERRORS"* ]]; then
        timestamp=$(echo "$line" | cut -d' ' -f1,2)
        device=$(echo "$line" | cut -d' ' -f4)
        io=$(echo "$line" | cut -d' ' -f5)
        sectors=$(echo "$line" | cut -d' ' -f6)
        ticks=$(echo "$line" | cut -d' ' -f7)
        ios=$(echo "$line" | cut -d' ' -f8)
        sectors_discard=$(echo "$line" | cut -d' ' -f9)
        ios_discard=$(echo "$line" | cut -d' ' -f10)
        sector_size=$(echo "$line" | cut -d' ' -f11)
        
        # 如果时间戳变化，添加一个空行
        if [[ -n "$current_timestamp" && "$current_timestamp" != "$timestamp" ]]; then
            echo "" >> "$OUTPUT_FILE"
        fi
        
        awk -v ts="$timestamp" -v d="$device" -v i="$io" -v s="$sectors" -v t="$ticks" -v id="$ios" -v sd="$sectors_discard" -v id="$ios_discard" -v ss="$sector_size" \
            'BEGIN {printf "%-20s %-8s %10d %10d %10d %10d %10d %10d %10d\n", ts, d, i, s, t, id, sd, id, ss}' >> "$OUTPUT_FILE"
        
        current_timestamp="$timestamp"
    fi
done < "$temp_disk_errors"

echo -e "\n=== 磁盘队列统计 ===" >> "$OUTPUT_FILE"
echo "时间戳                设备    最大请求数  当前请求数" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"
current_timestamp=""
while read -r line; do
    if [[ "$line" == *"DISK_QUEUE"* ]]; then
        timestamp=$(echo "$line" | cut -d' ' -f1,2)
        device=$(echo "$line" | cut -d' ' -f4)
        nr_requests=$(echo "$line" | cut -d' ' -f5)
        nr_in_flight=$(echo "$line" | cut -d' ' -f6)
        
        # 如果时间戳变化，添加一个空行
        if [[ -n "$current_timestamp" && "$current_timestamp" != "$timestamp" ]]; then
            echo "" >> "$OUTPUT_FILE"
        fi
        
        awk -v ts="$timestamp" -v d="$device" -v nr="$nr_requests" -v nf="$nr_in_flight" \
            'BEGIN {printf "%-20s %-8s %10d %10d\n", ts, d, nr, nf}' >> "$OUTPUT_FILE"
        
        current_timestamp="$timestamp"
    fi
done < "$temp_disk_queue"

echo -e "\n=== 文件系统使用统计 ===" >> "$OUTPUT_FILE"
echo "时间戳                设备    总大小  已用  可用  使用率  挂载点" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"
current_timestamp=""
while read -r line; do
    if [[ "$line" == *"FILESYSTEM_USAGE"* ]]; then
        timestamp=$(echo "$line" | cut -d' ' -f1,2)
        device=$(echo "$line" | cut -d' ' -f4)
        size=$(echo "$line" | cut -d' ' -f5)
        used=$(echo "$line" | cut -d' ' -f6)
        avail=$(echo "$line" | cut -d' ' -f7)
        usage=$(echo "$line" | cut -d' ' -f8)
        mount=$(echo "$line" | cut -d' ' -f9)
        
        # 如果时间戳变化，添加一个空行
        if [[ -n "$current_timestamp" && "$current_timestamp" != "$timestamp" ]]; then
            echo "" >> "$OUTPUT_FILE"
        fi
        
        awk -v ts="$timestamp" -v d="$device" -v s="$size" -v u="$used" -v a="$avail" -v us="$usage" -v m="$mount" \
            'BEGIN {printf "%-20s %-8s %-8s %-8s %-8s %-8s %s\n", ts, d, s, u, a, us, m}' >> "$OUTPUT_FILE"
        
        current_timestamp="$timestamp"
    fi
done < "$temp_filesystem_usage"

echo -e "\n=== I/O等待时间统计 ===" >> "$OUTPUT_FILE"
echo "时间戳                I/O等待时间(%)" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"
current_timestamp=""
while read -r line; do
    if [[ "$line" == *"IO_WAIT"* ]]; then
        timestamp=$(echo "$line" | cut -d' ' -f1,2)
        io_wait=$(echo "$line" | cut -d' ' -f4)
        
        # 如果时间戳变化，添加一个空行
        if [[ -n "$current_timestamp" && "$current_timestamp" != "$timestamp" ]]; then
            echo "" >> "$OUTPUT_FILE"
        fi
        
        awk -v ts="$timestamp" -v iw="$io_wait" \
            'BEGIN {printf "%-20s %10.2f%%\n", ts, iw}' >> "$OUTPUT_FILE"
        
        current_timestamp="$timestamp"
    fi
done < "$temp_io_wait"

# 计算平均值
echo -e "\n=== 统计摘要 ===" >> "$OUTPUT_FILE"
echo "------------------------------------------------" >> "$OUTPUT_FILE"

# 磁盘I/O平均值
for device in $(get_block_devices); do
    r_avg=$(awk -v d="$device" '/DISK_IO/ && $4==d {sum+=$7; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_io")
    w_avg=$(awk -v d="$device" '/DISK_IO/ && $4==d {sum+=$8; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_io")
    rsec_avg=$(awk -v d="$device" '/DISK_IO/ && $4==d {sum+=$9; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_io")
    wsec_avg=$(awk -v d="$device" '/DISK_IO/ && $4==d {sum+=$10; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_io")
    echo "设备 $device 平均I/O: 读请求=${r_avg}/s, 写请求=${w_avg}/s, 读扇区=${rsec_avg}/s, 写扇区=${wsec_avg}/s" >> "$OUTPUT_FILE"
done

# 磁盘利用率平均值
for device in $(get_block_devices); do
    util_avg=$(awk -v d="$device" '/DISK_UTIL/ && $4==d {sum+=$5; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_util")
    echo "设备 $device 平均利用率: ${util_avg}%" >> "$OUTPUT_FILE"
done

# 磁盘延迟平均值
for device in $(get_block_devices); do
    await_avg=$(awk -v d="$device" '/DISK_LATENCY/ && $4==d {sum+=$5; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_latency")
    svctm_avg=$(awk -v d="$device" '/DISK_LATENCY/ && $4==d {sum+=$6; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_latency")
    echo "设备 $device 平均延迟: 等待时间=${await_avg}ms, 服务时间=${svctm_avg}ms" >> "$OUTPUT_FILE"
done

# 磁盘吞吐量平均值
for device in $(get_block_devices); do
    tps_avg=$(awk -v d="$device" '/DISK_THROUGHPUT/ && $4==d {sum+=$5; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_throughput")
    kb_read_avg=$(awk -v d="$device" '/DISK_THROUGHPUT/ && $4==d {sum+=$6; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_throughput")
    kb_wrtn_avg=$(awk -v d="$device" '/DISK_THROUGHPUT/ && $4==d {sum+=$7; count++} END {print (count>0) ? sum/count : 0}' "$temp_disk_throughput")
    echo "设备 $device 平均吞吐量: ${tps_avg} IO/s, 读取=${kb_read_avg}KB/s, 写入=${kb_wrtn_avg}KB/s" >> "$OUTPUT_FILE"
done

# I/O等待时间平均值
io_wait_avg=$(awk '/IO_WAIT/ {sum+=$4; count++} END {print (count>0) ? sum/count : 0}' "$temp_io_wait")
echo "平均I/O等待时间: ${io_wait_avg}%" >> "$OUTPUT_FILE"

echo -e "\n结束时间: $(date)" >> "$OUTPUT_FILE"
echo "----------------------------------------" >> "$OUTPUT_FILE"

# 正常完成，调用cleanup但不退出
cleanup 
